package com.xshang.tenant;

import org.springframework.lang.NonNull;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * 租户数据源切换器：使用拦截器来实现切换
 */
public class TenantDataSourceSwitcher implements HandlerInterceptor {

    /**
     * 租户id请求头名称
     */
    private final String tenantIdHeaderName;

    @Resource
    private TenantDataSourceService tenantDataSourceService;

    /**
     * 租户数据源切换器缺省构造方法
     */
    public TenantDataSourceSwitcher() {
        this.tenantIdHeaderName = "tenant-id";
    }

    /**
     * 租户数据源切换器构造方法
     *
     * @param tenantIdHeaderName 租户id名称
     */
    public TenantDataSourceSwitcher(String tenantIdHeaderName) {
        this.tenantIdHeaderName = tenantIdHeaderName;
    }

    @Override
    public boolean preHandle(HttpServletRequest request,
                             @NonNull HttpServletResponse response,
                             @NonNull Object handler) {
        String tenantId = request.getHeader(tenantIdHeaderName);
        String tenantDataSourceName = TenantDataSourceContext.DEFAULT_TENANT_DATA_SOURCE;
        if (StringUtils.hasText(tenantId)) {
            // 数据库中查询租户对应的数据源名称，可以优化为从缓存中获取，提升性能
            String tenantDataSourceBeanName = tenantDataSourceService.getTenantDataSourceBeanName(tenantId.trim());
            if (StringUtils.hasText(tenantDataSourceBeanName)) {
                tenantDataSourceName = tenantDataSourceBeanName;
            }
        }
        // 设置上下文变量
        TenantDataSourceContext.set(tenantDataSourceName);
        return true;
    }

    @Override
    public void afterCompletion(@NonNull HttpServletRequest request,
                                @NonNull HttpServletResponse response,
                                @NonNull Object handler, Exception ex) {
        // 清理上下文变量
        TenantDataSourceContext.clear();
    }
}
